"""
AI Agent Pro - 数据库配置和连接管理
"""

import logging
from typing import AsyncGenerator

from sqlalchemy import MetaData, create_engine, event
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool

from app.core.config import get_settings

logger = logging.getLogger(__name__)
settings = get_settings()

# 创建异步数据库引擎
async_engine = create_async_engine(
    settings.database_url,
    echo=settings.debug,
    pool_pre_ping=True,
    pool_recycle=300,
    pool_size=20,
    max_overflow=30,
)

# 创建同步数据库引擎 (用于 Alembic 迁移)
sync_engine = create_engine(
    settings.database_url_sync,
    echo=settings.debug,
    pool_pre_ping=True,
    pool_recycle=300,
    pool_size=10,
    max_overflow=20,
)

# 创建会话工厂
AsyncSessionLocal = async_sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False,
    autocommit=False,
    autoflush=False,
)

SessionLocal = sessionmaker(
    sync_engine,
    autocommit=False,
    autoflush=False,
)

# 创建基础模型类
metadata = MetaData(
    naming_convention={
        "ix": "ix_%(column_0_label)s",
        "uq": "uq_%(table_name)s_%(column_0_name)s",
        "ck": "ck_%(table_name)s_%(constraint_name)s",
        "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
        "pk": "pk_%(table_name)s"
    }
)

Base = declarative_base(metadata=metadata)


async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """获取数据库会话 (异步)"""
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise
        finally:
            await session.close()


def get_sync_db():
    """获取数据库会话 (同步)"""
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


async def init_db():
    """初始化数据库"""
    try:
        # 测试数据库连接
        async with async_engine.begin() as conn:
            # 这里可以执行一些初始化 SQL
            await conn.execute("SELECT 1")
        
        logger.info("✅ 数据库连接成功")
        
        # 如果需要，可以在这里创建表
        # async with async_engine.begin() as conn:
        #     await conn.run_sync(Base.metadata.create_all)
        
    except Exception as e:
        logger.error(f"❌ 数据库连接失败: {e}")
        raise


async def close_db():
    """关闭数据库连接"""
    await async_engine.dispose()
    sync_engine.dispose()
    logger.info("🔒 数据库连接已关闭")


# 数据库事件监听器
@event.listens_for(sync_engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
    """设置 SQLite pragma (如果使用 SQLite)"""
    if "sqlite" in settings.database_url:
        cursor = dbapi_connection.cursor()
        cursor.execute("PRAGMA foreign_keys=ON")
        cursor.close()


# Redis 连接配置
import redis.asyncio as redis
from redis.asyncio import ConnectionPool

# 创建 Redis 连接池
redis_pool = ConnectionPool.from_url(
    settings.redis_url,
    encoding="utf-8",
    decode_responses=True,
    max_connections=20,
    retry_on_timeout=True,
)

# 创建 Redis 客户端
redis_client = redis.Redis(connection_pool=redis_pool)


async def get_redis():
    """获取 Redis 客户端"""
    return redis_client


async def init_redis():
    """初始化 Redis 连接"""
    try:
        await redis_client.ping()
        logger.info("✅ Redis 连接成功")
    except Exception as e:
        logger.error(f"❌ Redis 连接失败: {e}")
        raise


async def close_redis():
    """关闭 Redis 连接"""
    await redis_client.close()
    logger.info("🔒 Redis 连接已关闭")


# 向量数据库连接 (Qdrant)
from qdrant_client import AsyncQdrantClient
from qdrant_client.http import models

# 创建 Qdrant 客户端
qdrant_client = AsyncQdrantClient(
    url=settings.qdrant_url,
    api_key=settings.qdrant_api_key,
    timeout=60,
)


async def get_qdrant():
    """获取 Qdrant 客户端"""
    return qdrant_client


async def init_vector_db():
    """初始化向量数据库"""
    try:
        # 测试连接
        collections = await qdrant_client.get_collections()
        logger.info(f"✅ Qdrant 连接成功, 当前有 {len(collections.collections)} 个集合")
        
        # 创建默认集合 (如果不存在)
        collection_name = "agent_memory"
        try:
            await qdrant_client.get_collection(collection_name)
        except Exception:
            # 集合不存在，创建新集合
            await qdrant_client.create_collection(
                collection_name=collection_name,
                vectors_config=models.VectorParams(
                    size=1536,  # OpenAI embedding 维度
                    distance=models.Distance.COSINE
                ),
                optimizers_config=models.OptimizersConfigDiff(
                    memmap_threshold=20000,
                    max_segment_size=200000,
                ),
                hnsw_config=models.HnswConfigDiff(
                    m=16,
                    ef_construct=100,
                    full_scan_threshold=10000,
                )
            )
            logger.info(f"✅ 创建集合 '{collection_name}' 成功")
        
    except Exception as e:
        logger.error(f"❌ Qdrant 连接失败: {e}")
        raise


async def close_vector_db():
    """关闭向量数据库连接"""
    await qdrant_client.close()
    logger.info("🔒 Qdrant 连接已关闭")


# 数据库依赖注入
class DatabaseDependency:
    """数据库依赖类"""
    
    def __init__(self):
        self.db = None
        self.redis = None
        self.qdrant = None
    
    async def __aenter__(self):
        self.db = AsyncSessionLocal()
        self.redis = redis_client
        self.qdrant = qdrant_client
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.db:
            await self.db.close()


# 导出
__all__ = [
    "Base",
    "async_engine",
    "sync_engine",
    "AsyncSessionLocal",
    "SessionLocal",
    "get_db",
    "get_sync_db",
    "init_db",
    "close_db",
    "redis_client",
    "get_redis",
    "init_redis",
    "close_redis",
    "qdrant_client",
    "get_qdrant",
    "init_vector_db",
    "close_vector_db",
    "DatabaseDependency",
]